import math
import os
import sys
import time
from pathlib import Path
import numpy as np
import cv2
import torch
import ultralytics
from ultralytics import YOLO

def detect(img,count):
    results = model.predict(source=img,imgsz=imgsz,conf=conf_thres,iou=iou_thres,
                            max_det=max_det,save=save,save_txt=save_txt,save_conf=save_conf,save_crop=save_crop,
                            name=f"exchange_auto_red_Video_20250113103836696_{count}")
    dect = results[0].plot()
    cv2.imshow("a",dect)
    cv2.waitKey(1)
    return dect

weights = r"/home/champrin/Desktop/rm24_arm_ws/src/arm_auto_exchange/exchange_slot_detector/model/weights/best.pt"

Video_or_ImagesPath = r"/home/champrin/Desktop/MV-CS016-10UC+DA1041860/a/exchange_blue_Video_20250113103836696_335.jpg"

imgsz = 640  # 输入图片的大小 默认640(pixels)
conf_thres = 0.4  # object置信度阈值 默认0.25  用在nms中
iou_thres = 0.4  # 做nms的iou阈值 默认0.45   用在nms中
max_det = 12  # 每张图片最多的目标数量  用在nms中

save = True
save_txt = False
save_conf = False
save_crop = False

if save == False:
    save_txt = False
    save_conf = False
# 载入模型
model = YOLO(weights)  # load FP32 model
count = 1
detect(Video_or_ImagesPath,count)




